package yycore

import chisel3._
import chisel3.util._
import chisel3.util.experimental.BoringUtils
import common.Constants._

trait BPUConst {
  val btbWays = 2

  val padLen = 2
  val NRbtb = 128
  val NRras = 16
  val NRbht = 64
  val bhr_width = log2Ceil(NRbht)

  val idxBits = log2Up(NRbtb)
  val tagBits = AddrBits - padLen - idxBits

  def BTBEntry() = new Bundle {
    val valid = Bool()
    val tag = UInt(tagBits.W)
    val btb_type = UInt(2.W)
    val target = UInt(AddrBits.W)
  }
  def btbAddrBundle() = new Bundle {
    val tag = UInt(tagBits.W)
    val index = UInt(idxBits.W)
    val pad = UInt(padLen.W)
  }

  def getBTBTag(pc: UInt) = pc.asTypeOf(btbAddrBundle()).tag
  def getBTBIdx(pc: UInt) = pc.asTypeOf(btbAddrBundle()).index
  def Hash(pc: UInt) = { pc }

}

object BTBtype {
  def B = "b00".U  // branch
  def J = "b01".U  // jump
  def I = "b10".U  // indirect
  def R = "b11".U  // return

  def apply() = UInt(2.W)
}

class BPUUpdateReq extends Bundle {
  val valid = Output(Bool())
  val pc = Output(UInt(AddrBits.W))
  val isMissPredict = Output(Bool())
  val actualTarget = Output(UInt(AddrBits.W))
  val actualTaken = Output(Bool())  // for branch
  val brType = Output(BRType())
  val btbType = Output(BTBtype())
}

class BPUOutIO extends Bundle {
  val target = Output(UInt(AddrBits.W))
  //val rtype = Output(UInt(1.W)) // 1: branch mispredict: only need to flush frontend  0: others: flush the whole pipeline
  val valid = Output(Bool())
}

/**
 * BPU
 */
class BPU extends Module with BPUConst{
  val io = IO(new Bundle(){
    val in = Flipped(ValidIO(new Bundle {
      val pc = Output(UInt(AddrBits.W))
    }))
    val out = new BPUOutIO
  })

  val pc = io.in.bits.pc
  val pcLatch = RegEnable(io.in.bits.pc, io.in.valid)
  /**
   * BTH Entry
   */
  val btbArray = SyncReadMem(NRbtb, BTBEntry())
  val btbRead = btbArray.read(getBTBIdx(pc), io.in.valid)
  val btbHit = btbRead.valid && btbRead.tag === getBTBTag(pcLatch) && RegNext(io.in.valid)
  when(reset.asBool()) {
    for (i <- 0 until NRbtb) {
      btbArray.write(i.U, 0.U.asTypeOf(BTBEntry()))
    }
  }

  /**
   * RAS
   */
  val ras = Mem(NRras, UInt(AddrBits.W))
  val sp = Counter(NRras)
  val rasTarget = RegEnable(next = ras(sp.value), init = 0.U, enable = io.in.valid)
  //val rasTarget = RegEnable(ras.read(sp.value), io.in.valid)
  when(reset.asBool()) {
    for (i <- 0 until NRras) {
      ras.write(i.U, 0.U)
    }
  }

  /**
   *  Branch History Table
   */
  def getBHTAddr(pc: UInt) = pc(bhr_width + 1, 2)
  val BHT = Mem(NRbht, UInt(bhr_width.W))
  val readBHR = BHT.read(getBHTAddr(pc))
  val phtAddr = getBHTAddr(pc) ^ readBHR
  when(reset.asBool()) {
    for (i <- 0 until NRbht) {
      BHT.write(i.U, 0.U)
    }
  }


  /**
   * PHT:00, 01, 10, 11初始化：weakly not taken
   * strongly not taken, weakly not taken, weakly taken, strongle taken
   */
  def isTaken(phr: UInt): Bool = {
    phr(1)
  }
  val pht = Mem(NRbht, UInt(2.W))
  val phtTaken = RegEnable(next = isTaken(pht(phtAddr)), init = false.B, enable = io.in.valid)
  //val phtTaken = RegEnable(pht.read(phtAddr)(1), io.in.valid)
  when(reset.asBool()) {
    for (i <- 0 until NRbht) {
      pht.write(i.U, 1.U)
    }
  }



  /**
   * misspredirect update BHT&PHT and BTB&RAS
   */

  val req = WireInit(0.U.asTypeOf(new BPUUpdateReq))
  BoringUtils.addSink(req, "bpuUpdateReq")

  // update BHT
  when (req.valid) {
    when (req.actualTaken) {
      val bhrOld = BHT(getBHTAddr(req.pc))
      val bhrNew = Mux(req.actualTaken, Cat(bhrOld(bhr_width - 1, 1), "b1".U), Cat(bhrOld(bhr_width - 1, 1), "b0".U))
      BHT.write(getBHTAddr(req.pc), bhrNew)
    }
  }

  // update PHT
  val phr = RegNext(pht.read(BHT.read(getBHTAddr(pc)) ^ getBHTAddr(pc)))
  val phr_addr = BHT.read(getBHTAddr(pc)) ^ getBHTAddr(pc)
  when (req.valid && BRType.isBranch(req.brType)) {
    val taken = req.actualTaken
    val newPhr = Mux(taken, phr + 1.U, phr - 1.U)
    val wen = (taken && (phr =/= "b11".U)) || (!taken && (phr =/= "b00".U))
    when (wen) {
      pht.write(phr_addr, newPhr)
    }
  }
  // update BTB
  val btbWrite = WireInit(0.U.asTypeOf(BTBEntry()))
  btbWrite.tag := getBTBTag(req.pc)
  btbWrite.target := req.actualTarget
  btbWrite.btb_type := req.btbType
  btbWrite.valid := true.B
  when(req.valid && req.isMissPredict ) {
    btbArray.write(getBTBIdx(req.pc), btbWrite)
  }
  // update RAS
  when (req.valid) {
    when (req.brType === BRType.call)  {
      ras.write(sp.value + 1.U, req.pc + 4.U)
      sp.value := sp.value + 1.U
    }
    .elsewhen (req.brType === BRType.ret) {
      when(sp.value === 0.U) {
        //val rasFailed = Wire(Bool())
      }
      sp.value := Mux(sp.value === 0.U, 0.U, sp.value - 1.U)
    }
  }

  /**
   * Output
   */
  //val pcLatch = RegNext(pc)
  io.out.valid := btbHit && Mux(btbRead.btb_type === BTBtype.B, phtTaken, rasTarget =/= 0.U)
  io.out.target := Mux(btbRead.btb_type === BTBtype.R /*&& !rasFailed*/, rasTarget, btbRead.target)
                   // Mux(btbHit, btbRead.target, pcLatch + 4.U))
}

